package uk.ac.cam.djjb2.mygame;

import java.util.AbstractCollection;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Queue;
import java.util.Set;

public class PriorityQueueSet<E> extends AbstractCollection<E> implements Queue<E>, Set<E> {
	private Map<E, Integer> lookup;
	private ArrayList<E> heap;
	private Comparator<E> comparator;
	
	public PriorityQueueSet() {
		lookup = new HashMap<E, Integer>();
		heap = new ArrayList<E>();
		comparator = null;
	}
	public PriorityQueueSet(Comparator<E> comparator) {
		lookup = new HashMap<E, Integer>();
		heap = new ArrayList<E>();
		this.comparator = comparator;
	}
	public PriorityQueueSet(Map<E, Integer> map) {
		lookup = map;
		heap = new ArrayList<E>();
		comparator = null;
	}
	public PriorityQueueSet(Map<E, Integer> map, Comparator<E> comparator) {
		lookup = map;
		heap = new ArrayList<E>();
		this.comparator = comparator;
	}

	private int getParent(int index) {
		return (index-1) >> 1;
	}
	private int getLeftChild(int index) {
		return (index<<1) + 1;
	}
	private int getMinChild(int index) {
		int child1 = getLeftChild(index);
		int child2 = child1 + 1;
		
		if (child2 >= heap.size())
			return child1;
		
		if (compareTo(child1, child2) <= 0)
			return child1;
		else
			return child2;
	}
	
	@SuppressWarnings("unchecked")
	private int compareTo(int a, int b) {
		if (comparator != null)
			return comparator.compare(heap.get(a), heap.get(b));
		else
			return ((Comparable<E>)heap.get(a)).compareTo(heap.get(b));
	}
	private void swap(int a, int b) {
		E tmp = heap.get(a);
		heap.set(a, heap.get(b));
		heap.set(b, tmp);
	}
	private void updateMap(int index) {
		lookup.put(heap.get(index), index);
	}
	private int increasePriority(int index) {
		int parent = getParent(index);
		while ((index > 0) && (compareTo(index, parent) < 0)) {
			swap(index, parent);
			updateMap(index);
			index = parent;
			parent = getParent(index);
		}
		return index;
	}
	private int decreasePriority(int index) {
		int child = getMinChild(index);
		while ((child < heap.size()) && (compareTo(index, child) > 0)) {
			swap(index, child);
			updateMap(index);
			index = child;
			child = getMinChild(index);
		}
		return index;
	}
	private E remove(int index) {
		E result = heap.get(index);
		
		if (index < heap.size()-1) {
			swap(index, heap.size()-1);
			lookup.remove(result);
			heap.remove(heap.size()-1);
			updateMap(decreasePriority(index));
		} else {
			lookup.remove(result);
			heap.remove(heap.size()-1);
		}
		
		return result;
	}
	
	public boolean add(E elem) {
		Integer index = lookup.get(elem);
		
		if (index == null) {
			heap.add(elem);
			updateMap(increasePriority(heap.size()-1));
			
			return true;
		}
		
		return false;
	}
	@Override
	public boolean remove(Object elem) {
		Integer index = lookup.get(elem);
		if (index != null) {
			remove((int)index);
			
			return true;
		}
		return false;
	}
	public void update(E elem) {
		Integer index = lookup.get(elem);
		if (index != null) {
			int i = increasePriority(index);
			if (i == index)
				i = decreasePriority(i);
			updateMap(i);
		}
	}
	public E peek() {
		if (!heap.isEmpty())
			return heap.get(0);
		else
			return null;
	}
	public E poll() {
		if (!heap.isEmpty())
			return remove(0);
		else
			return null;
	}
	public int size() {
		return heap.size();
	}
	@Override
	public boolean isEmpty() {
		return heap.isEmpty();
	}
	@Override
	public boolean contains(Object paramObject) {
		return lookup.containsKey(paramObject);
	}
	@Override
	public Iterator<E> iterator() {
		return new Iterator<E>() {
			int index = heap.size();

			@Override
			public boolean hasNext() {
				return index > 0;
			}

			@Override
			public E next() {
				if (index > 0)
					return heap.get(--index);
				else
					throw new NoSuchElementException();
			}

			@Override
			public void remove() {
				heap.remove(index);
			}
		};
	}
	@Override
	public Object[] toArray() {
		return heap.toArray();
	}
	@Override
	public <T> T[] toArray(T[] paramArrayOfT) {
		return heap.toArray(paramArrayOfT);
	}
	@Override
	public boolean containsAll(Collection<?> paramCollection) {
		return lookup.keySet().containsAll(paramCollection);
	}
	@Override
	public boolean addAll(Collection<? extends E> paramCollection) {
		boolean result = false;
		
		for (E elem : paramCollection)
			result |= add(elem);
		
		return result;
	}
	@Override
	public boolean removeAll(Collection<?> paramCollection) {
		boolean result = false;
		
		for (Object elem : paramCollection)
			result |= remove(elem);
		
		return result;
	}
	@Override
	public boolean retainAll(Collection<?> paramCollection) {
		boolean result = false;
		
		for (int i = heap.size()-1; i >= 0; i--)
			if (!paramCollection.contains(heap.get(i))) {
				result |= true;
				remove(i);
			}
		
		return result;
	}
	@Override
	public void clear() {
		for (int i = 0; i < heap.size(); i++)
			lookup.remove(heap.get(i));
		heap.clear();
	}
	@Override
	public boolean offer(E paramE) {
		return add(paramE);
	}
	@Override
	public E remove() {
		return poll();
	}
	@Override
	public E element() {
		return peek();
	}
}
